
import numpy as np
import tensorflow as tf
from keras.callbacks import TensorBoard
from keras.layers import Input, Dense
from keras.models import Model
 
import keras
from keras.datasets import mnist
from keras.models import Sequential
from keras.layers import Dense, Dropout
from keras.optimizers import RMSprop
from keras.utils import plot_model
batch_size = 128
num_classes = 10
epochs = 20

def write_log(callback, names, logs, batch_no):
    for name, value in zip(names, logs):
        summary = tf.Summary()
        summary_value = summary.value.add()
        summary_value.simple_value = value
        summary_value.tag = name
        callback.writer.add_summary(summary, batch_no)
        callback.writer.flush()

model = Sequential()
model.add(Dense(512, activation='relu', input_shape=(784,)))
model.add(Dropout(0.2))
model.add(Dense(512, activation='relu'))
model.add(Dropout(0.2))
model.add(Dense(num_classes, activation='softmax'))
model.compile(loss='categorical_crossentropy',
              optimizer=RMSprop(),
              metrics=['accuracy'])


(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train = x_train.reshape(60000, 784)
x_test = x_test.reshape(10000, 784)
x_train = x_train.astype('float32')
x_test = x_test.astype('float32')
x_train /= 255
x_test /= 255


log_path = './logs'
callback = TensorBoard(log_path)
callback.set_model(model)
train_names = ['train_loss', 'train_mae']
val_names = ['val_loss', 'val_mae']
for batch_no in range(1000):
    logs = model.train_on_batch(x_train, y_train)
    write_log(callback, train_names, logs, batch_no)






# net_in = Input(shape=(3,))
# net_out = Dense(1)(net_in)
# model = Model(net_in, net_out)
# model.compile(loss='mse', optimizer='sgd', metrics=['mae'])
 
# log_path = './logs'
# callback = TensorBoard(log_path)
# callback.set_model(model)
# train_names = ['train_loss', 'train_mae']
# val_names = ['val_loss', 'val_mae']
# for batch_no in range(100):
#     X_train, Y_train = np.random.rand(32, 3), np.random.rand(32, 1)
#     logs = model.train_on_batch(X_train, Y_train)
#     write_log(callback, train_names, logs, batch_no)
    
#     # if batch_no % 10 == 0:
#     #     X_val, Y_val = np.random.rand(32, 3), np.random.rand(32, 1)
#     #     logs = model.train_on_batch(X_val, Y_val)
#     #     write_log(callback, val_names, logs, batch_no//10)
